# import libs
import os
import sys
import numpy as np
import torch
import torchvision
import torchvision.transforms as tT
from PIL import Image, ExifTags
import json

#检查标注文件的正确性，把错误标注写入err_files.txt
def check_annotations(folder:str):
    # check the annotaion files and write checking results to the local file ann_err.txt

    files = list(sorted(os.listdir(folder)))
    im_files = [f for f in files if f.endswith("jpg")]
    print("总共有："+str(len(im_files))+"张jpg格式的图像，应该是水表图像。")

    no_jpg_json_files = []
    no_png_json_files = []
    no_png_mask_files = []

    with open(os.path.join(folder, "err_files.txt"), 'w') as err_file:
        for f in im_files:
            err_lst = [] #保存针对该图像文件的错误信息
            # 检查是否缺基本标注文件和掩膜标注或缺掩膜图像文件
            if f+".json" not in files:
                no_jpg_json_files.append(f)
                print(f+"缺少对应的标注文件")
                err_lst.append("缺少对应的标注文件")
                if len(no_jpg_json_files)==1:
                    print("存在缺少基本标注json的文件")           

            else:

                # if f[:-4]+".png" not in files:
                #     no_png_mask_files.append(f)
                #     err_lst.append("缺少对应的轮廓掩膜Png图像")
                #     print(f + "缺少对应的轮廓掩膜Png图像")
                #     if len(no_png_mask_files)==1:
                #         print("\n存在缺少轮廓掩膜png的文件")

                # if f[:-4]+".png"+".json" not in files:
                #     err_lst.append("缺少轮廓掩膜json的文件")
                #     no_png_json_files.append(f)
                #     if len(no_png_json_files)==1:
                #         print("\n存在缺少轮廓掩膜json的文件")


                #读取基本标记和掩膜信息，并检查细节的合法性    
                file_path = os.path.join(folder, f+".json")
                with open(file_path, 'r', encoding='utf-8') as fjson:
                    injson = json.load(fjson)

                    #3针的标注方式与8针的不同，多了一个step，估计是表盘轮廓
                if_4_steps = "step_4" in injson.keys()
                if if_4_steps:
                    key_bboxes = "step_2"
                    key_zero_scale_dir = "step_3"
                    key_pointers_dir = "step_4"
                else:
                    key_bboxes = "step_1"
                    key_zero_scale_dir = "step_2"
                    key_pointers_dir = "step_3"
                    
                    #读取指针的包围盒
                n_pointers = len(injson[key_bboxes]['result'])
                if (n_pointers!=3 and n_pointers!=8 and n_pointers!=4):
                    err_lst.append("指针数目量不是3也不是8")    

                    #读取指针方向                
                n_pointers2 = len(injson[key_pointers_dir]['result'])
                if(n_pointers!=n_pointers2):
                    err_lst.append("包围盒的数量与指针方向标注的数量不一致")

                    #检查指针方向的两点是否距离过近                      
                for line in injson[key_pointers_dir]['result']:
                    pt_lst = line['pointList']      
                    line_vec = np.array([pt_lst[0]['x']-pt_lst[1]['x'], pt_lst[0]['y']-pt_lst[1]['y']] )                
                    if np.linalg.norm(line_vec)<20:
                        err_lst.append("指针方向的两点距离过近")
                        break

                    #检查指针是否在包围盒内     
                    # ！！注意：一定要先讲包围盒列表和指针标注列表按id进行排序，避免没有对齐的问题。 
                injson[key_bboxes]['result'].sort(key=lambda x:x["id"])   
                injson[key_pointers_dir]['result'].sort(key=lambda x:x["sourceID"])   
                if(n_pointers==n_pointers2):
                    for i in range(n_pointers):
                        bbox = injson[key_bboxes]['result'][i] # dict
                        pointer = injson[key_pointers_dir]['result'][i]['pointList'] # dict
                        if_in_bbox = False
                        for j in range(2):
                            pt = pointer[j]
                            if 0<pt['x']-bbox['x']<=bbox['width'] and 0<pt['y']-bbox['y']<=bbox['height']:
                                if_in_bbox = True
                                break
                        if if_in_bbox == False:
                            err_lst.append("指针方向的标注没有位于对应的包围盒内")
                            break


                    #检查零刻度方向的两点是否距离过近
                n_zero_scale_dir = len(injson[key_zero_scale_dir]['result'])            
                for line in injson[key_zero_scale_dir]['result']:
                    pt_lst = line['pointList']      
                    line_vec = np.array([pt_lst[0]['x']-pt_lst[1]['x'], pt_lst[0]['y']-pt_lst[1]['y']] )
                    if np.linalg.norm(line_vec)<20:
                        err_lst.append("零刻度方向的两点距离过近")
                        break

                    #zero_scale_dir_lst.append([pt_lst[0]['x'],pt_lst[0]['y'],pt_lst[1]['x'],pt_lst[1]['y']]) # [x1,y1,x2,y2]第1个点是起点,据说是从刻度0到5


            #如有错误，则把该图像所包含的错误信息写入文件
            if len(err_lst)>0:
                err_file.write('\n')
                err_file.write(f+":")
                for err in err_lst:
                    err_file.write("||"+err)

    print("\n检查annotation完毕")

def load_valid_imfile_names(folder:str)->list:
    files = list(sorted(os.listdir(folder)))
    im_files = [f for f in files if f.endswith("jpg")]
    #filter out the illegal files 
    assert(os.path.exists(os.path.join(folder, "err_files.txt")))
    illegel_lst = []
    with open(os.path.join(folder, "err_files.txt"),'r',errors='ignore') as err_file:
        lines = err_file.readlines()
        for line in lines:
            if len(line)>0:
                pos_comma=line.find(':')
                if pos_comma>-1:
                    illegel_lst.append(line[:pos_comma])
    
    legal_lst = []
    for f in im_files:
        if f not in illegel_lst:
            legal_lst.append(f)
    
    return legal_lst

def load_anno(folder:str, im_fname:str, down_scale_factor=8)->dict:
    #读取标注的json文件内容
    file_path = os.path.join(folder, im_fname+".json")
    with open(file_path, 'r', encoding='utf-8') as f:
        injson = json.load(f)
    #3针的标注方式与8针的不同，多了一个step，估计是表盘轮廓,8针就没有key_pointer_mask
    if_4_steps = "step_4" in injson.keys()
    key_pointer_mask = None
    if if_4_steps:
        if injson["step_4"]["toolName"]=="polygonTool":
            key_bboxes = "step_1"
            key_zero_scale_dir = "step_2"
            key_pointers_dir = "step_3"
            key_pointer_mask = "step_4"
        else:
            key_bboxes = "step_2"
            key_zero_scale_dir = "step_3"
            key_pointers_dir = "step_4"
    else:
        key_bboxes = "step_1"
        key_zero_scale_dir = "step_2"
        key_pointers_dir = "step_3"
    
    # ！！注意：一定要先讲包围盒列表和指针标注列表按id进行排序，避免没有对齐的问题。 
    injson[key_bboxes]['result'].sort(key=lambda x:x["id"])   
    injson[key_pointers_dir]['result'].sort(key=lambda x:x["sourceID"])   #它的sourceID对应包围盒的id

    #读取包围盒信息(貌似有点重复了，感觉就是表盘轮廓包围盒)
    n_pointers = len(injson[key_bboxes]['result'])
    bboxes_lst = []    
    for bbox in injson[key_bboxes]['result']:
        bboxes_lst.append([bbox['x']//down_scale_factor,bbox['y']//down_scale_factor,bbox['width']//down_scale_factor,bbox['height']//down_scale_factor])        
    
    #读取零刻度方向  
    n_zero_scale_dir = len(injson[key_zero_scale_dir]['result'])
    zero_scale_dir_lst = []
    for line in injson[key_zero_scale_dir]['result']:
        pt_lst = line['pointList']        
        zero_scale_dir_lst.append([pt_lst[0]['x']//down_scale_factor,pt_lst[0]['y']//down_scale_factor,pt_lst[1]['x']//down_scale_factor,pt_lst[1]['y']//down_scale_factor]) # [x1,y1,x2,y2]第1个点是起点,据说是从刻度0到5
        
    #读取指针方向
    n_pointers2 = len(injson[key_pointers_dir]['result'])   
    pointer_dir_lst = [] 
    for line in injson[key_pointers_dir]['result']:
        pt_lst = line['pointList']
        if n_pointers2==8:
            pointer_dir_lst.append([pt_lst[0]['x']//down_scale_factor,pt_lst[0]['y']//down_scale_factor,pt_lst[1]['x']//down_scale_factor,pt_lst[1]['y']//down_scale_factor]) # [x1,y1,x2,y2]第1个点是起点，是旋转中心
        elif n_pointers2==3:
            pointer_dir_lst.append([pt_lst[1]['x']//down_scale_factor,pt_lst[1]['y']//down_scale_factor,pt_lst[0]['x']//down_scale_factor,pt_lst[0]['y']//down_scale_factor]) # [x1,y1,x2,y2]第1个点是起点，是旋转中心
    

    '''下面这段根据情况看是否需要注释'''
    #读取表盘轮廓  
    mask_fname = im_fname[:-4]+".png"
    mask = Image.open(os.path.join(folder, mask_fname))
    # mask = rotate_im_accord_exiftag(mask)
    im_sz=(mask.size[0]//down_scale_factor, mask.size[1]//down_scale_factor)
    mask = mask.resize(im_sz)
    mask = np.array(mask)
    if len(mask.shape)==3:
        mask = np.max(mask, axis=2) #每个图中只有一个instance，而每个图是RGBA格式的，所以将其压缩成1个通道
    mask = (mask>0).astype(np.int8)*1#转成由0和1组成的图像
    obj_ids = np.unique(mask)
    obj_ids = obj_ids[1:]#去除0这个代表背景的标签
    masks = mask == obj_ids[:, None, None]


    '''下面这段根据情况看是否需要注释'''
    #得到表盘轮廓对应的包围盒
    num_meters = len(obj_ids)
    mboxes = []
    for i in range(num_meters):
        pos = np.where(masks[i])
        xmin = np.min(pos[1])
        xmax = np.max(pos[1])
        ymin = np.min(pos[0])
        ymax = np.max(pos[0])
        mboxes.append([xmin, ymin, xmax, ymax])    
    
    #得到指针轮廓(8针的key_pointer_mask=None)
    # n_black_pointer = len(injson[key_pointer_mask]['result'])
    # pointer_polygon_lst = []
    # for i in range(n_pointers):
    #     pointer_id = injson[key_bboxes]['result'][i]['id']
    #     polygon = None
    #     for j in range(n_black_pointer):
    #         if injson[key_pointer_mask]['result'][j]['sourceID']==pointer_id:
    #             ptlist = injson[key_pointer_mask]['result'][j]['pointList']
    #             n_pt = len(ptlist)
    #             polygon = np.zeros((n_pt,2))
    #             for k in range(n_pt):
    #                 polygon[k,:][0] = ptlist[k]["x"]//down_scale_factor
    #                 polygon[k,:][1] = ptlist[k]["y"]//down_scale_factor

    #             break
    #     pointer_polygon_lst.append(polygon)


    anno={"bbox":bboxes_lst, "zdir":zero_scale_dir_lst, "pdir":pointer_dir_lst, "mask_boxes":mboxes,"masks":masks,}#, "masks":masks, "mask_boxes":mboxes, "ppoly":pointer_polygon_lst}
    return anno

def rotate_im_accord_exiftag(im:Image)->Image:
    #处理手机拍照时的旋转问题(ref: https://blog.csdn.net/qq_34904125/article/details/120918232)
    for orientation in ExifTags.TAGS.keys():
        if ExifTags.TAGS[orientation]=='Orientation':
            break
    exif = dict(im._getexif().items())
    if orientation in exif:       
        if exif[orientation] == 3:
            im = im.rotate(180,expand=True)
        elif exif[orientation]==6:
            im = im.rotate(270,expand=True)
        elif exif[orientation]==8:
            im=im.rotate(90, expand=True)
    
    return im

class MeterPtDirDataset(torch.utils.data.Dataset):
    def __init__(self, root, transforms, down_scale_factor=8):
        self.root = root # root folder path
        self.transforms = transforms # data transformations
        self.down_scale_factor = down_scale_factor # 图像缩小为原图的比例
                            
        self.imgs = load_valid_imfile_names(root)
    
    def __len__(self):
        return len(self.imgs)
    
    def __getitem__(self, idx):        
        fpath = os.path.join(self.root, self.imgs[idx])
        img = Image.open(fpath)
        img = rotate_im_accord_exiftag(img)
        img = img.convert('RGB')
        #缩小图像
        im_sz=(img.size[0]//self.down_scale_factor, img.size[1]//self.down_scale_factor)
        img = img.resize(im_sz)
        
        #载入标注信息
        # anno = load_anno(self.root, self.imgs[idx], self.down_scale_factor)
        anno = load_anno(self.root, self.imgs[idx], 1)
        #获得包围盒信息
        boxes = []
        for b in anno['bbox']:
            #缩小原始的数值
            for i in range(4):
                b[i] = b[i]/self.down_scale_factor
            
            #转换为x0,y0,x1,y1格式
            bbox = [b[0],b[1],b[0]+b[2],b[1]+b[3]]
            boxes.append(bbox)
        
        boxes = torch.as_tensor(boxes, dtype=torch.float32)

        #获得keypoint信息
        keypoints = []
        for dir in anno['pdir']:
            pt1 = [v/self.down_scale_factor for v in dir[:2]]
            pt1.append(1)
            pt2 = [v/self.down_scale_factor for v in dir[2:]]
            pt2.append(1)
            pts = [pt1, pt2]

            keypoints.append(pts)
        
        keypoints = torch.as_tensor(keypoints, dtype=torch.float32)

        #获得类别信息，只有1类目标就是仪表
        num_objs = len(boxes)
        labels = torch.ones((num_objs,), dtype=torch.int64)
        
        #获取其他coco格式的必要信息
        image_id = torch.tensor([idx])
        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
        iscrowd = torch.zeros((num_objs,), dtype=torch.int64)

        #获得标签信息
        target = {}
        target["boxes"] = boxes
        target["keypoints"] = keypoints
        target["labels"] = labels        
        target["image_id"] = image_id
        target["area"] = area
        target["iscrowd"] = iscrowd

        #进行变化
        if self.transforms is not None:
            img, target = self.transforms(img, target)
        
        return img, target        


#展示图片
def show(imgs):
        if not isinstance(imgs, list):
            imgs = [imgs]
        fig, axs = plt.subplots(ncols=len(imgs), squeeze=False)
        for i, img in enumerate(imgs):
            img = img.detach()
            img = F.to_pil_image(img)
            axs[0, i].imshow(np.asarray(img))
            axs[0, i].set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])



def iou(x1, y1, w1, h1, x2, y2, w2, h2):
    # 计算两个矩形的面积
    area1 = w1 * h1
    area2 = w2 * h2

    # 计算两个矩形的交集面积
    x_left = max(x1, x2)
    y_top = max(y1, y2)
    x_right = min(x1 + w1, x2 + w2)
    y_bottom = min(y1 + h1, y2 + h2)
    if x_right < x_left or y_bottom < y_top:
        intersection = 0
    else:
        intersection = (x_right - x_left) * (y_bottom - y_top)

    # 计算两个矩形的并集面积
    union = area1 + area2 - intersection

    # 计算IoU
    iou = intersection / union

    return iou


# 我自己写的NUM函数，goal_num是要选取的目标个数8个框
def my_NMS(boxes,scores,threshold,goal_num):
    len_s=len(scores)
    # if goal_num<0 or goal_num>len_s:#貌似有时候会检测不到8个
    if goal_num<0:
        return False
    used,select_idx=[0]*len_s,[]
    for i in range(goal_num):
        max_score,max_score_idx=-1,-1
        for j in range(len_s):#找到第i轮最大的score
            if used[j]==1:
                continue
            if scores[j]>max_score:
                max_score=scores[j]
                max_score_idx=j
        ''' TO do'''
        if select_idx==-1:#找完了所有的
            break
        select_idx.append(max_score_idx)
        used[max_score_idx]=1
        for j in range(len_s):
            if used[j]==1:
                continue
            box=boxes[j,:]
            s_box=boxes[max_score_idx,:]
            if iou(box[0],box[1],box[2]-box[0],box[3]-box[1],s_box[0],s_box[1],s_box[2]-s_box[0],s_box[3]-s_box[1])>threshold:
                used[j]=1
    return select_idx


'''新的去除置信得分低的预测，方法实现'''
def remove_low_scores(boxes,scores,threshold):
    # 获取所有大于等于 0.9 的元素的索引
    idx = scores >= 0.9
    # 使用索引操作获取对应的元素
    scores = scores[idx]
    boxes=boxes[idx]
    return boxes,scores



# 定义拟合圆的函数
from scipy.optimize import leastsq
def fit_circle(x, y):
    # 定义误差函数
    def err(p, x, y):
        return (x - p[0])**2 + (y - p[1])**2 - p[2]**2
    # 初始圆心和半径的估计值
    p0 = [np.mean(x), np.mean(y), np.std(x)]
    # 使用最小二乘法求解圆心和半径
    p, success = leastsq(err, p0, args=(x, y))
    # 返回圆心和半径
    return p[:2], p[2]



#定义获得关键点，返回以拟合圆的中心划分的顺时针方向的关键点
def get_center_seq(keypoints,selected_idx=[0,1,2,3,4,5,6,7]):
    #分别获取nms后,输入有selected_idx则为输入的，否则默认前置信度前八的8个点的x和y坐标
    x,y=[],[]
    for i in selected_idx:
        # print(f'i:{i}')
        x.append(keypoints[i,0,0])
        y.append(keypoints[i,0,1])
    # 测试拟合圆的函数
    # print("x:",x)
    # print("Y:",y)
    center, radius = fit_circle(x, y)
    # print(center,radius)
    # return [center,radius]
    #以圆心y坐标为界划分为2部分
    up,down=[],[]
    d_idx,u_idx=[],[]
    for i in selected_idx:
        if keypoints[i,0,1]>center[1]:#y比圆心y大则在我们正常视角的下方
            down.append([keypoints[i,0,0],keypoints[i,0,1],keypoints[i,1,0],keypoints[i,1,1]])
            d_idx.append(i)
        else:
            up.append([keypoints[i,0,0],keypoints[i,0,1],keypoints[i,1,0],keypoints[i,1,1]])
            u_idx.append(i)

    # print("len_UP:",len(up),"len_DOWN:",len(down))
    # print("up_index:",u_idx,"down_index:",d_idx)

    #up按照x从小到大，down按照x从大到小，以此实现顺时针
    up.sort()
    down.sort(reverse=True)
    total=up+down
    return total



def get_center_seq_blur(keypoints,selected_idx=[0,1,2,3,4,5,6,7]):
    #分别获取8个点的x和y坐标
    x,y=[],[]
    for i in selected_idx:
        x.append(keypoints[i,0,0])
        y.append(keypoints[i,0,1])
    # 测试拟合圆的函数
    # print("x:",x)
    # print("Y:",y)
    center, radius = fit_circle(x, y)
    # print(center,radius)
    # return [center,radius]
    #以圆心x坐标为界划分为2部分
    left,right=[],[]
    l_idx,r_idx=[],[]
    for i in selected_idx:
        if keypoints[i,0,0]>center[0]:#x比圆心x大则在右方
            right.append([keypoints[i,0,0],keypoints[i,0,1],keypoints[i,1,0],keypoints[i,1,1]])
            r_idx.append(i)
        else:
            left.append([keypoints[i,0,0],keypoints[i,0,1],keypoints[i,1,0],keypoints[i,1,1]])
            r_idx.append(i)

    # print("len_UP:",len(up),"len_DOWN:",len(down))
    # print("up_index:",u_idx,"down_index:",d_idx)

    #left按照y从小到大，right按照y从大到小，以此实现顺时针
    left= sorted(left, key=lambda x: x[1])
    right = sorted(right, key=lambda x: x[1], reverse=True)
    total=left+right
    return total




'''比较Image-Based Automatic Dial Meter Reading in Unconstrained Scenarios中的相邻读数纠正方法'''
def judge(num1,num2): #我的
    a,b=int(num1),int(num2)
    x=(num1-a)*10#x为num1的小数部分*10
    threshold=6
    if x>b and x-b>=threshold:
        return (a+1)%10
    elif b>x and b-x>=threshold:
        return (a-1+10)%10
    else:
        return a
def judge2(num1,num2):#论文的
    a,b=int(num1),int(num2)
    x,y=(num1-a)*10,(num2-b)*10#x，y为num1，num2的小数部分*10
    if x>7 and y<6:
        return (a+1)%10
    else:
        return a
def get_reading(x_o,y_o,x_p,y_p,x_l,y_l):#o为圆点，p为指针尖(默认ol方向为与o水平向右)
    o=(x_o,y_o)
    l=(x_l,y_l)
    p=(x_p,y_p)
    # 向量AB和向量CD的坐标表示
    ol = np.array([l[0] - o[0], l[1] - o[1]])
    op = np.array([p[0] - o[0], p[1] - o[1]])

    # 计算向量AB和向量CD的夹角
    angle = np.arctan2(ol[0]*op[1]-ol[1]*op[0], ol[0]*op[0]+ol[1]*op[1])
    angle = angle * 180 / np.pi  # 将弧度转换为角度


    if angle < 0:#逆时针方向的angle是负，+360转换为顺时针方向的角度
        angle += 360
    # 输出夹角
    reading=int((angle/360)*100)/10#先乘100再除10为了得到保留一位小数的读数
    return reading

def get_reading_zdir(x_o,y_o,x_p,y_p,zdir):#o为圆点，p为指针尖,zdir中是x1，y1,x2,y2（依次5和0刻度)
    x_l,y_l,x_l2,y_l2=zdir
    o=(x_o,y_o)
    l=(x_l,y_l)

    p=(x_p,y_p)
    l_2=(x_l2,y_l2)
    # 向量AB和向量CD的坐标表示

    l_l_2 = np.array([l_2[0] - l[0], l_2[1] - l[1]])
    op = np.array([p[0] - o[0], p[1] - o[1]])

    # 计算向量AB和向量CD的夹角
    angle = np.arctan2(l_l_2[0]*op[1]-l_l_2[1]*op[0], l_l_2[0]*op[0]+l_l_2[1]*op[1])
    angle = angle * 180 / np.pi  # 将弧度转换为角度


    if angle < 0:#逆时针方向的angle是负，+360转换为顺时针方向的角度
        angle += 360
    # 输出夹角
    reading=int((angle/360)*100)/10#先乘100再除10为了得到保留一位小数的读数
    return reading

def pointer_to_read(total,num,zdir=None):
    total_reading,times,end_reading_before,end_reading_after=[],1000,0,0
    if zdir==None:#位指定零刻度方向，则默认水平方向       
        for i in range(num):
            zdir=[total[i][0],total[i][1],total[i][0]+1,total[i][1]]
            # reading=get_reading(total[i][0],total[i][1],total[i][2],total[i][3],total[0][0]+1,total[0][1])
            reading=get_reading_zdir(total[i][0],total[i][1],total[i][2],total[i][3],zdir)
            total_reading.append(reading)
    else:
        for i in range(num):
            reading=get_reading_zdir(total[i][0],total[i][1],total[i][2],total[i][3],zdir)
            total_reading.append(reading)

    print("带小数点：",total_reading)

    # for i in range(num):
    #     # total_reading[i]=int(total_reading[i])#直接取整
    #     if i!=num-1:
    #         # total_reading[i]=judge(total_reading[i],total_reading[i+1])#我的方法纠正
    #         total_reading[i]=judge2(total_reading[i],total_reading[i+1])
    #     else:
    #         total_reading[i]=int(total_reading[i])
    for i in range(num-1,-1,-1):#从后往前修正
        # total_reading[i]=int(total_reading[i])#直接取整
        if i!=num-1:
            total_reading[i]=judge(total_reading[i],total_reading[i+1])#我的方法纠正
            # total_reading[i]=judge2(total_reading[i],total_reading[i+1])#论文方法纠正
        else:
            total_reading[i]=int(total_reading[i])
    return total_reading
